//
//  NSData+CryptorAdditions.m
//  DaumComics
//
//  Created by tiburon on 11. 1. 11..
//  Copyright 2011 DaumCommunications. All rights reserved.
//
#import <CommonCrypto/CommonCryptor.h>
#import <zlib.h>

#import "NSData+CryptorAdditions.h"

@interface NSData (Utils)
- (void)logError:(CCCryptorStatus)cryptStatus;
@end


@implementation NSData (CryptorAdditions)
//////////////////////////////////////////////////////////////////////////////////////////////////
- (NSData *)DESEncryptWithKeyString:(NSString*)keyString ccOptions:(CCOptions)ccOptions{
	char keyPtr[kCCKeySizeDES+1]; // room for terminator (unused)
	bzero(keyPtr, sizeof(keyPtr)); // fill with zeroes (for padding)
	
	// fetch key data
	[keyString getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
	
	NSUInteger dataLength = [self length];
	
	//See the doc: For block ciphers, the output size will always be less than or 
	//equal to the input size plus the size of one block.
	//That's why we need to add the size of one block here
	size_t bufferSize = dataLength + kCCBlockSizeAES128;
	void *buffer = malloc(bufferSize);
	
	size_t numBytesEncrypted = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt, kCCAlgorithmDES, kCCOptionPKCS7Padding,
										  keyPtr, kCCKeySizeDES,
										  NULL /* initialization vector (optional) */,
										  [self bytes], dataLength, /* input */
										  buffer, bufferSize, /* output */
										  &numBytesEncrypted);
	if (cryptStatus == kCCSuccess) {
		//the returned NSData takes ownership of the buffer and will free it on deallocation
		return [NSData dataWithBytesNoCopy:buffer length:numBytesEncrypted];
	}
	
	free(buffer); //free the buffer;
	return nil;
}
//////////////////////////////////////////////////////////////////////////////////////////////////
- (NSData *)DESDecryptWithKeyString:(NSString*)keyString ccOptions:(CCOptions)ccOptions{
	char keyPtr[kCCKeySizeDES+1]; // room for terminator (unused)
	bzero(keyPtr, sizeof(keyPtr)); // fill with zeroes (for padding)
	
	// fetch key data
	[keyString getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
	
	NSUInteger dataLength = [self length];
	
	//See the doc: For block ciphers, the output size will always be less than or 
	//equal to the input size plus the size of one block.
	//That's why we need to add the size of one block here
	size_t bufferSize = dataLength + kCCBlockSizeAES128;
	void *buffer = malloc(bufferSize);
	
	size_t numBytesEncrypted = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt, kCCAlgorithmDES, kCCOptionPKCS7Padding,
										  keyPtr, kCCKeySizeDES,
										  NULL /* initialization vector (optional) */,
										  [self bytes], dataLength, /* input */
										  buffer, bufferSize, /* output */
										  &numBytesEncrypted);
	if (cryptStatus == kCCSuccess) {
		//the returned NSData takes ownership of the buffer and will free it on deallocation
		return [NSData dataWithBytesNoCopy:buffer length:numBytesEncrypted];
	}
	
	free(buffer); //free the buffer;
	return nil;
}


- (NSData *)DESEncryptWithKeyData:(NSData*)keyData ccOptions:(CCOptions)ccOptions{
	
	
	NSUInteger dataLength = [self length];
	
	//See the doc: For block ciphers, the output size will always be less than or 
	//equal to the input size plus the size of one block.
	//That's why we need to add the size of one block here
	size_t bufferSize = dataLength + kCCBlockSizeDES;
	void *buffer = malloc(bufferSize);
	
	size_t numBytesEncrypted = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt, kCCAlgorithmDES, ccOptions,
										  [keyData bytes], kCCKeySizeDES,
										  NULL /* initialization vector (optional) */,
										  [self bytes], dataLength, /* input */
										  buffer, bufferSize, /* output */
										  &numBytesEncrypted);
	if (cryptStatus == kCCSuccess) {
		//the returned NSData takes ownership of the buffer and will free it on deallocation
		return [NSData dataWithBytesNoCopy:buffer length:numBytesEncrypted];
	}
	
	free(buffer); //free the buffer;
	return nil;
}
- (NSData *)DESDecryptWithKeyData:(NSData*)keyData ccOptions:(CCOptions)ccOptions{
	NSUInteger dataLength = [self length];
	
	//See the doc: For block ciphers, the output size will always be less than or 
	//equal to the input size plus the size of one block.
	//That's why we need to add the size of one block here
	size_t bufferSize = dataLength + kCCBlockSizeDES;
	void *buffer = malloc(bufferSize);
	
	size_t numBytesDecrypted = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt, kCCAlgorithmDES, ccOptions,
										  [keyData bytes], kCCKeySizeDES,
										  NULL /* initialization vector (optional) */,
										  [self bytes], dataLength, /* input */
										  buffer, bufferSize, /* output */
										  &numBytesDecrypted);
	
	if (cryptStatus == kCCSuccess) {
		//the returned NSData takes ownership of the buffer and will free it on deallocation
		return [NSData dataWithBytesNoCopy:buffer length:numBytesDecrypted];
	}
	
	free(buffer); //free the buffer;
	return nil;
    
}
//
////////////////////////////////////////////////////////////////////////////////////////////////////
//- (NSData *)DESEncryptWithKeyData:(NSData*)keyData ccOptions:(CCOptions)ccOptions{
//	unsigned char *input = (unsigned char*)[self bytes];
//	NSUInteger inLength = [self length];
//	NSInteger outLength = ((inLength + kCCBlockSizeDES) & ~(kCCBlockSizeDES - 1));
//	
//	unsigned char *output =(unsigned char *)calloc(outLength, sizeof(unsigned char));
//	bzero(output, outLength*sizeof(unsigned char));
//	size_t additionalNeeded = 0;
//	CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt, 
//										  kCCAlgorithmDES, 
//										  ccOptions, 
//										  [keyData bytes], 
//										  kCCKeySizeDES, 
//										  NULL, 
//										  input, 
//										  inLength, 
//										  output, 
//										  outLength, 
//										  &additionalNeeded); 
//	if(cryptStatus!=kCCSuccess){
//		NSLog(@"DES encrypt err: %d", cryptStatus);
//		[self logError:cryptStatus];	
//	}
////	if([self length]==8){
////		for(int i=0; i<[self length]; i++){
////			NSLog(@"===output %x", output[i]);
////		}
////	}
//	return [NSData dataWithBytesNoCopy:output length:outLength freeWhenDone:YES];
//}
//
////////////////////////////////////////////////////////////////////////////////////////////////////
//- (NSData *)DESDecryptWithKeyData:(NSData*)keyData ccOptions:(CCOptions)ccOptions{
//	unsigned char *input = (unsigned char*)[self bytes];
//	NSUInteger inLength = [self length];
//	NSInteger outLength = ((inLength + kCCBlockSizeDES) & ~(kCCBlockSizeDES - 1));
//	
//	unsigned char *output =(unsigned char *)calloc(outLength, sizeof(unsigned char));
//	bzero(output, outLength*sizeof(unsigned char));
//	size_t additionalNeeded = 0;
//	
//	CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt, 
//										  kCCAlgorithmDES, 
//										  ccOptions, 
//										  [keyData bytes], 
//										  kCCKeySizeDES, 
//										  NULL, 
//										  input, 
//										  inLength, 
//										  output, 
//										  outLength, 
//										  &additionalNeeded); 	
//	if(cryptStatus!=kCCSuccess){
//		NSLog(@"DES decrypt err: %d", cryptStatus);
//		[self logError:cryptStatus];	
//	} 
////	if([self length]==8){
////		for(int i=0; i<[self length]; i++){
////			NSLog(@"===output %x", output[i]);
////		}
////	}
//	return [NSData dataWithBytesNoCopy:output length:outLength freeWhenDone:YES];
//}
//


//////////////////////////////////////////////////////////////////////////////////////////////////
- (NSData *)AESEncryptWithKeyString:(NSString*)keyString ccOptions:(CCOptions)ccOptions{
	char keyPtr[kCCKeySizeAES256+1]; // room for terminator (unused)
	bzero(keyPtr, sizeof(keyPtr)); // fill with zeroes (for padding)
	[keyString getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
	
	
	unsigned char *input = (unsigned char*)[self bytes];
	NSUInteger inLength = [self length];
	NSInteger outLength = ((inLength + kCCBlockSizeAES128) & ~(kCCBlockSizeAES128 - 1));
	
	unsigned char *output =(unsigned char *)calloc(outLength, sizeof(unsigned char));
	bzero(output, outLength*sizeof(unsigned char));
	size_t additionalNeeded = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt, 
										  kCCAlgorithmAES128, 
										  ccOptions, 
										  keyPtr, 
										  kCCKeySizeAES256, 
										  NULL, 
										  input, 
										  inLength, 
										  output, 
										  outLength, 
										  &additionalNeeded); 
	if(cryptStatus!=kCCSuccess){
		NSLog(@"AES encrypt err: %d", cryptStatus);
		[self logError:cryptStatus];	
	}
	return [NSData dataWithBytesNoCopy:output length:outLength freeWhenDone:YES];
}

//////////////////////////////////////////////////////////////////////////////////////////////////
- (NSData *)AESDecryptWithKeyString:(NSString*)keyString ccOptions:(CCOptions)ccOptions{
	char keyPtr[kCCKeySizeAES256+1]; // room for terminator (unused)
	bzero(keyPtr, sizeof(keyPtr)); // fill with zeroes (for padding)
	[keyString getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
	
	
	unsigned char *input = (unsigned char*)[self bytes];
	NSUInteger inLength = [self length];
	NSInteger outLength = ((inLength + kCCBlockSizeAES128) & ~(kCCBlockSizeAES128 - 1));
	
	unsigned char *output =(unsigned char *)calloc(outLength, sizeof(unsigned char));
	bzero(output, outLength*sizeof(unsigned char));
	size_t additionalNeeded = 0;
	
	CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt, 
										  kCCAlgorithmAES128, 
										  ccOptions, 
										  keyPtr, 
										  kCCKeySizeAES256, 										  
										  NULL, 
										  input, 
										  inLength, 
										  output, 
										  outLength, 
										  &additionalNeeded); 	
	if(cryptStatus!=kCCSuccess){
		NSLog(@"AES decrypt err: %d", cryptStatus);
		[self logError:cryptStatus];	
	} 
	if([self length]==8){
		for(int i=0; i<[self length]; i++){
			NSLog(@"===output %x", output[i]);
		}
	}
	return [NSData dataWithBytesNoCopy:output length:outLength freeWhenDone:YES];
}
//////////////////////////////////////////////////////////////////////////////////////////////////
- (NSData *)AESEncryptWithKeyData:(NSData*)keyData ccOptions:(CCOptions)ccOptions{
	unsigned char *input = (unsigned char*)[self bytes];
	NSUInteger inLength = [self length];
	NSInteger outLength = ((inLength + kCCBlockSizeAES128) & ~(kCCBlockSizeAES128 - 1));
	
	unsigned char *output =(unsigned char *)calloc(outLength, sizeof(unsigned char));
	bzero(output, outLength*sizeof(unsigned char));
	size_t additionalNeeded = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt, 
										  kCCAlgorithmAES128, 
										  ccOptions, 
										  [keyData bytes], 
										  kCCKeySizeAES256, 
										  NULL, 
										  input, 
										  inLength, 
										  output, 
										  outLength, 
										  &additionalNeeded); 
	if(cryptStatus!=kCCSuccess){
		NSLog(@"AES encrypt err: %d", cryptStatus);
		[self logError:cryptStatus];	
	}
	return [NSData dataWithBytesNoCopy:output length:outLength freeWhenDone:YES];
}

//////////////////////////////////////////////////////////////////////////////////////////////////
- (NSData *)AESDecryptWithKeyData:(NSData*)keyData ccOptions:(CCOptions)ccOptions{
	unsigned char *input = (unsigned char*)[self bytes];
	NSUInteger inLength = [self length];
	NSInteger outLength = ((inLength + kCCBlockSizeAES128) & ~(kCCBlockSizeAES128 - 1));
	
	unsigned char *output =(unsigned char *)calloc(outLength, sizeof(unsigned char));
	bzero(output, outLength*sizeof(unsigned char));
	size_t additionalNeeded = 0;
	
	CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt, 
										  kCCAlgorithmAES128, 
										  ccOptions, 
										  [keyData bytes], 
										  kCCKeySizeAES256, 										  
										  NULL, 
										  input, 
										  inLength, 
										  output, 
										  outLength, 
										  &additionalNeeded); 	
	if(cryptStatus!=kCCSuccess){
		NSLog(@"AES decrypt err: %d", cryptStatus);
		[self logError:cryptStatus];	
	} 
	if([self length]==8){
		for(int i=0; i<[self length]; i++){
			NSLog(@"===output %x", output[i]);
		}
	}
	return [NSData dataWithBytesNoCopy:output length:outLength freeWhenDone:YES];
}

- (NSData *)AES256EncryptWithKeyData:(NSString *)key {
	// 'key' should be 32 bytes for AES256, will be null-padded otherwise
	char keyPtr[kCCKeySizeAES256+1]; // room for terminator (unused)
	bzero(keyPtr, sizeof(keyPtr)); // fill with zeroes (for padding)
	
	// fetch key data
	[key getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
	
	NSUInteger dataLength = [self length];
	
	//See the doc: For block ciphers, the output size will always be less than or 
	//equal to the input size plus the size of one block.
	//That's why we need to add the size of one block here
	size_t bufferSize = dataLength + kCCBlockSizeAES128;
	void *buffer = malloc(bufferSize);
	
	size_t numBytesEncrypted = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt, kCCAlgorithmAES128, kCCOptionPKCS7Padding,
										  keyPtr, kCCKeySizeAES256,
										  NULL /* initialization vector (optional) */,
										  [self bytes], dataLength, /* input */
										  buffer, bufferSize, /* output */
										  &numBytesEncrypted);
	if (cryptStatus == kCCSuccess) {
		//the returned NSData takes ownership of the buffer and will free it on deallocation
		return [NSData dataWithBytesNoCopy:buffer length:numBytesEncrypted];
	}
	
	free(buffer); //free the buffer;
	return nil;
}

- (NSData *)AES256DecryptWithKeyData:(NSString *)key {
	// 'key' should be 32 bytes for AES256, will be null-padded otherwise
	char keyPtr[kCCKeySizeAES256+1]; // room for terminator (unused)
	bzero(keyPtr, sizeof(keyPtr)); // fill with zeroes (for padding)
	
	// fetch key data
	[key getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
	
	NSUInteger dataLength = [self length];
	
	//See the doc: For block ciphers, the output size will always be less than or 
	//equal to the input size plus the size of one block.
	//That's why we need to add the size of one block here
	size_t bufferSize = dataLength + kCCBlockSizeAES128;
	void *buffer = malloc(bufferSize);
	
	size_t numBytesDecrypted = 0;
	CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt, kCCAlgorithmAES128, kCCOptionPKCS7Padding,
										  keyPtr, kCCKeySizeAES256,
										  NULL /* initialization vector (optional) */,
										  [self bytes], dataLength, /* input */
										  buffer, bufferSize, /* output */
										  &numBytesDecrypted);
	
	if (cryptStatus == kCCSuccess) {
		//the returned NSData takes ownership of the buffer and will free it on deallocation
		return [NSData dataWithBytesNoCopy:buffer length:numBytesDecrypted];
	}
	
	free(buffer); //free the buffer;
	return nil;
}


#pragma mark Utility methods
//////////////////////////////////////////////////////////////////////////////////////////////////
- (void)logError:(CCCryptorStatus)cryptStatus{
	if (cryptStatus == kCCParamError) NSLog(@"PARAM ERROR");
	else if (cryptStatus == kCCBufferTooSmall) NSLog(@"BUFFER TOO SMALL");
	else if (cryptStatus == kCCMemoryFailure) NSLog(@"MEMORY FAILURE");
	else if (cryptStatus == kCCAlignmentError) NSLog(@"ALIGNMENT");
	else if (cryptStatus == kCCDecodeError) NSLog(@"DECODE ERROR");
	else if (cryptStatus == kCCUnimplemented) NSLog(@"UNIMPLEMENTED");	
}
@end

